from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Conv1D, Flatten, Dropout
from tensorflow.keras.optimizers import Adam
from scikeras.wrappers import KerasClassifier
from sklearn.base import BaseEstimator, ClassifierMixin


class CustomKerasCNNClassifier(BaseEstimator, ClassifierMixin):
    def __init__(self, filters=64, kernel_size=3, dense_units=64, learning_rate=0.001, epochs=10, batch_size=32):
        self.filters = filters
        self.kernel_size = kernel_size
        self.dense_units = dense_units
        self.learning_rate = learning_rate
        self.epochs = epochs
        self.batch_size = batch_size
        self.model_ = None

    def _build_model(self):
        model = Sequential([
            Conv1D(filters=self.filters,
                   kernel_size=self.kernel_size,
                   activation='relu',
                   input_shape=self.input_shape_),
            Flatten(),
            Dense(self.dense_units, activation='relu'),
            Dropout(0.5),
            Dense(1, activation='sigmoid')
        ])
        model.compile(
            loss='binary_crossentropy',
            optimizer=Adam(learning_rate=self.learning_rate),
            metrics=['accuracy']
        )
        return model

    def fit(self, X, y=None, **kwargs):
        self.input_shape_ = X.shape[1:]
        self.model_ = self._build_model()
        self.model_.fit(X, y, epochs=self.epochs, batch_size=self.batch_size, verbose=0)
        return self

    def predict(self, X):
        preds = (self.model_.predict(X, verbose=0) > 0.5).astype(int).flatten()
        return preds

    def predict_proba(self, X):
        probs = self.model_.predict(X, verbose=0).flatten()
        return probs

    def get_params(self, deep=True):
        return {
            "filters": self.filters,
            "kernel_size": self.kernel_size,
            "dense_units": self.dense_units,
            "learning_rate": self.learning_rate,
            "epochs": self.epochs,
            "batch_size": self.batch_size
        }

    def set_params(self, **parameters):
        for parameter, value in parameters.items():
            setattr(self, parameter, value)
        return self
